import torch 
from dataset import SliceBuilder
# import nrrd 
# import SimpleITK as sitk
from dataset import StandardHDF5Dataset
import dataset 
import matplotlib.pyplot as plt 
import matplotlib
# print(matplotlib.get_backend())
import h5py
import numpy as np

from skimage.segmentation import find_boundaries

def _recover_ignore_index(input, orig, ignore_index):
    if ignore_index is not None:
        mask = orig == ignore_index
        input[mask] = ignore_index

    return input

def blur_boundary(boundary, sigma):
    boundary = gaussian(boundary, sigma=sigma)
    boundary[boundary >= 0.5] = 1
    boundary[boundary < 0.5] = 0
    return boundary

class StandardLabelToBoundary:
    def __init__(self, ignore_index=None, append_label=False, blur=False, sigma=1, mode='thick', foreground=False,
                 **kwargs):
        self.ignore_index = ignore_index
        self.append_label = append_label
        self.blur = blur
        self.sigma = sigma
        self.mode = mode
        self.foreground = foreground

    def __call__(self, m):
        assert m.ndim == 3

        boundaries = find_boundaries(m, connectivity=2, mode=self.mode)
        if self.blur:
            boundaries = blur_boundary(boundaries, self.sigma)

        results = []
        if self.foreground:
            foreground = (m > 0).astype('uint8')
            results.append(_recover_ignore_index(foreground, m, self.ignore_index))

        results.append(_recover_ignore_index(boundaries, m, self.ignore_index))

        if self.append_label:
            # append original input data
            results.append(m)

        return np.stack(results, axis=0)



if __name__ == "__main__":
    # import torch.nn as nn 
    # t1 = torch.rand((5, 1, 32, 32, 32))
    # transConv = nn.ConvTranspose3d(1, 1, kernel_size=2, stride=2,
    #                                            padding=0)
    # print(transConv(t1).shape)
    ## 测试下3d 残差网络输入输出
    # from unet3d.model import ResidualUNet3D
    # t1 = torch.rand((5, 1, 32, 32, 32))
    # net = ResidualUNet3D(in_channels=1, out_channels=1, final_sigmoid=True) # 训练阶段 testing = False 最后一层不会使用sigmoid进行激活。

    # print(net(t1).shape)

    # import os 
    # if not os.path.exists("./checkpointss/"):
    #     os.mkdir("./checkpointss/")

    
    # datasets = StandardHDF5Dataset.create_datasets(["./test_h5/"])

    # print(datasets)
    #### ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 测试一下 get_train_loaders 这个函数有没有问题，应该是没问题。
    # loder = dataset.get_train_loaders("./test_h5/", "./test_h5/")
    # train_loder = loder["train"]
    # for src, tgt in train_loder:
    #     print(src.shape)
    #     print(tgt.shape)
    #     plt.imshow(src.squeeze(0).squeeze(0)[10], cmap="gray")
    #     plt.show()
    #     plt.imshow(tgt.squeeze(0).squeeze(0)[10], cmap="gray")
    #     plt.show()
    #～～～～～～～～～～～～～～～～～～～～～～～～ 测试下高亮边界预测这个数据集。
    # import h5py
    # # from unet3d.metrics import BoundaryAdaptedRandError
    # f = h5py.File('./light_boundary_out/out.h5','r')   #打开h5文件
    # # print(f.keys())                          #可以查看所有的主键
    # # # a = f['data'][:]                    #取出主键为data的所有的键值
    # # # a = f["raw"][()][100, :, :]
    # # # b = f["label"][()][100, :, :]

    # # for i in range(5, 30):
    # res = torch.tensor(f["predictions"][()])[0, 200]
    # plt.imshow(res, cmap="gray")
    # plt.show()

    # f2 = h5py.File("./lightsheet_boundary_val/Movie1_t00004_crop_gt.h5")
    # res = f2["label"][()]
    # print(res.shape)
    # plt.imshow(res[200], cmap="gray")
    # plt.show()

    # label2 = f["label"][()]
    # lb = StandardLabelToBoundary()
    # label2 = lb(label2)
    # plt.imshow(label2[0][100], cmap="gray")
    # plt.show()
    # # print(res)
    # print(label2.shape)
    # bar = BoundaryAdaptedRandError(thresholds=[0.4])
    # seg = bar.input_to_segm(label2[0][100])
    # print(seg.shape)
    # plt.imshow(seg, cmap="gray")
    # plt.show()
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~``

    # print(slice_bulid.raw_slices)

    # nrrd_data, nrrd_options = nrrd.read("./meningioma_data/DUAN_YU_SHU/PN0224273_1.2.840.113564.9435757939512.9868.636587794821385442.378/1.3.12.2.1107.5.2.19.46085.2018040814385458714883968.0.0.0/T1WI-CE_t_1.nrrd")

    # print(nrrd_data.shape)

    # reader = sitk.ImageSeriesReader()
    # dicom_image_names = reader.GetGDCMSeriesFileNames("./meningioma_data/DUAN_YU_SHU")
    # reader.SetFileNames(dicom_image_names)
    # image_agency = reader.Execute()

    # # print(image_agency)
    # image_volume = sitk.GetArrayFromImage(image_agency)
    # print(image_volume.shape)


    # ### ~~~~~~~~~测试一下3d 细胞分割数据
    # import h5py
    # from augment.transforms import BlobsToMask

    # # from unet3d.metrics import BoundaryAdaptedRandError
    # f = h5py.File('./cell_seg_3d_train/t00015_s01_uint8_cropped_gt.h5','r')   #打开h5文件
    # # print(f.keys())                          #可以查看所有的主键
    # # # a = f['data'][:]                    #取出主键为data的所有的键值
    # # # a = f["raw"][()][100, :, :]
    # # # b = f["label"][()][100, :, :]

    # # for i in range(5, 30):
    # bm = BlobsToMask()
    # print(f.keys())
    # res = f["raw"][()]
    # print(res.shape)
    # res1 = f["label"][()]
    # res1 = bm(res1)
    # print(res1.shape)
    # plt.imshow(res[200], cmap="gray")
    # plt.show()
    # plt.imshow(res1[0][200], cmap="gray")
    # plt.show()

    # # f2 = h5py.File("./lightsheet_boundary_val/Movie1_t00004_crop_gt.h5")
    # # res = f2["label"][()]
    # # print(res.shape)
    # # plt.imshow(res[200], cmap="gray")
    # # plt.show()


    #~~~~~~~~~~~~~~~~~测试下自己的数据

    # f = h5py.File('./self_data_val/brain_data18t1.h5','r')   #打开h5文件
    # print(f.keys())                          #可以查看所有的主键
    # # a = f['data'][:]                    #取出主键为data的所有的键值
    # a = f["raw"][()]
    # b = f["label"][()]

    # plt.imshow(a[17], cmap="gray")
    # plt.show()
    # plt.imshow(b[17], cmap="gray")
    # plt.show()

    # print(a.shape)
    # print(b.shape)


    f = h5py.File('./self_data_val/brain_data21fl.h5','r')   #打开h5文件
    print(f.keys())                          #可以查看所有的主键
    # a = f['data'][:]                    #取出主键为data的所有的键值
    a = f["raw"][()]
    b = f["label"][()]

    print(a.shape)
    print(b.shape)

    # plt.imshow(a[13], cmap="gray")
    # plt.show()
    # plt.imshow(b[13], cmap="gray")
    # plt.show()
    # print(a.shape)
    # print(b.shape)

    # f = h5py.File('./self_out0.h5','r')   #打开h5文件
    # print(f.keys())                          #可以查看所有的主键
    # # a = f['data'][:]                    #取出主键为data的所有的键值
    # a = f["predictions"][()]
    # print(a.shape)

    # plt.imshow(a[0][13], cmap="gray")
    # plt.show()


